package com.ttai.untils;

import com.ttai.untils.SRPUtil;

import java.math.BigInteger;
import java.security.MessageDigest;

/**
 * @program: taisecurity
 * @description: SRP客户端逻辑
 * @author: Jingxi Zhu
 * @create: 2018-11-20 20:55
 * @Copyright © 2018 Jingxi Zhu. All rights reserved
 **/
public class SRPClient {

    private BigInteger n;
    private BigInteger g;
    private byte[] s;
    private BigInteger x;
    private BigInteger v;
    private BigInteger a;
    private BigInteger A;
    private String user;
    private byte[] key;
    private MessageDigest hash, ckhash;

    private static int A_LEN = 64;		// 64 bits for 'a'

    /**
     * Creates a new SRP Client object from the initial round of the
     * exchange.
     * @param username The user's username on the server host.
     * @param modulus The user's safe-prime modulus, received from
     *                the server.
     * @param generator The user's primitive generator, received from
     *                  the server.
     * @param salt The user's password salt, received from the server.
     */
    public SRPClient(String username, byte[] modulus, byte[] generator,
                     byte[] salt) {
        user = username;
        n = new BigInteger(1, modulus);
        g = new BigInteger(1, generator);
        s = salt;
        key = null;
        hash = SRPUtil.newDigest();
        hash.update(SRPUtil.xor(SRPUtil.newDigest().digest(modulus),
                SRPUtil.newDigest().digest(generator), 20));
        hash.update(SRPUtil.newDigest().digest(username.getBytes()));
        hash.update(salt);
        ckhash = SRPUtil.newDigest();
    }

    /**
     * @returns The user's safe-prime modulus
     */
    public byte[] modulus() { return SRPUtil.trim(n.toByteArray()); }

    /**
     * @returns The user's primitive generator
     */
    public byte[] generator() { return SRPUtil.trim(g.toByteArray()); }

    /**
     * @returns The user's password salt
     */
    public byte[] salt() { return s; }

    /**
     * @returns The exponential residue (parameter A) to be sent to the
     *          server.
     */
    public byte[] exponential() {
        if(A == null) {
            BigInteger one = BigInteger.valueOf(1);
            do {
                a = new BigInteger(A_LEN, SRPUtil.RNG);
            } while(a.compareTo(one) <= 0);
            A = g.modPow(a, n);
            byte[] out = SRPUtil.trim(A.toByteArray());
            hash.update(out);
            ckhash.update(out);
            return out;
        }
        else{
            return SRPUtil.trim(A.toByteArray());
        }

    }

    /**
     * Deprecated.  Use exponential() instead.
     */
    public byte[] generateExponential() { return exponential(); }

    /**
     * Incorporates the user's password into the session key computation.
     * @param pass The user's password or passphrase.
     */
    public void inputPassword(String pass) {
        MessageDigest ctxt = SRPUtil.newDigest();
        ctxt.update(s);
        ctxt.update(SRPUtil.userHash(user, pass));
        x = new BigInteger(1, ctxt.digest());
        v = g.modPow(x, n);
    }

    /**
     * @returns The secret shared session key between client and server
     * @param srvexp The server's exponential (parameter B).
     */
    public byte[] sessionKey(byte[] srvexp) {
        hash.update(srvexp);
        byte[] uhash = SRPUtil.newDigest().digest(srvexp);
        byte[] fourbytes = {uhash[0], uhash[1], uhash[2], uhash[3]};
        BigInteger sum = x.multiply(new BigInteger(1, fourbytes));
        BigInteger base = new BigInteger(1, srvexp);
        if(base.compareTo(v) < 0){
            base = base.add(n);
        }

        base = base.subtract(v);
        BigInteger S = base.modPow(sum.add(a), n);
        key = SRPUtil.sessionKeyHash(SRPUtil.trim(S.toByteArray()));
        hash.update(key);
        return key;
    }

    /**
     * @returns The secret shared session key between client and server
     */
    public byte[] sessionKey() { return key; }

    /**
     * Deprecated.  Use sessionKey() instead.
     */
    public byte[] getSessionKey(byte[] srvexp) {
        return sessionKey(srvexp);
    }

    /**
     * @returns The response to the server's challenge.
     */
    public byte[] response() {
        byte[] resp = hash.digest();
        ckhash.update(resp);
        ckhash.update(key);
        return resp;
    }

    /**
     * @param resp The server's response to the client's challenge
     * @returns True if and only if the server's response was correct.
     */
    public boolean verify(byte[] resp) {
        return SRPUtil.matches(resp, ckhash.digest());
    }

}
